import pandas as pd
from pathlib import Path
import openpyxl
from openpyxl.styles import Font, numbers
from typing import Dict

from interfaces.excel_writer import IExcelWriter


class ExcelWriter(IExcelWriter):
    def __init__(self, output_path: str, column_delete_list: list, get_column_data_fn):
        self.output_path = Path(output_path)
        self.column_delete_list = column_delete_list
        self.get_column_data_fn = get_column_data_fn

    def _clean_sheet_name(self, sheet_name: str) -> str:
        invalid_chars = '\\/:*?[]'
        for char in invalid_chars:
            sheet_name = sheet_name.replace(char, '_')
        return sheet_name[:31]

    def _set_sheet_text_format(self, ws):
        for row in ws.iter_rows():
            for cell in row:
                cell.number_format = numbers.FORMAT_TEXT

    def _enable_sheet_filter(self, ws):
        if ws.max_row > 1 and ws.max_column > 0:
            ws.auto_filter.ref = f"A1:{openpyxl.utils.get_column_letter(ws.max_column)}{ws.max_row}"

    def _set_sheet_font(self, ws):
        times_new_roman_font = Font(name='Times New Roman', size=11)
        for row in ws.iter_rows():
            for cell in row:
                cell.font = times_new_roman_font

    def _remove_columns(self, all_results: Dict[str, pd.DataFrame]) -> Dict[str, pd.DataFrame]:
        for key, data in list(all_results.items()):
            if isinstance(data, pd.DataFrame) and not data.empty:
                columns_to_keep = [col for col in data.columns if col not in self.column_delete_list]
                all_results[key] = data[columns_to_keep]
        return all_results

    def _reorder_dict_unspecified_first(self, all_results: Dict[str, pd.DataFrame], key_order: list) -> Dict[str, pd.DataFrame]:
        all_keys = list(all_results.keys())
        # 优先写入的表名（若存在）：确保它们排在最前
        priority_first = ['入组成功受试者', '失败受试者列表', '实验室关键指标']
        priority_keys = [k for k in priority_first if k in all_results]

        specified_keys = [key for key in key_order if key in all_results]
        priority_and_specified = set(priority_keys) | set(specified_keys)
        # 其余未指定的表按自然插入顺序，其后再写入指定表
        unspecified_keys = [key for key in all_keys if key not in priority_and_specified]

        new_key_order = priority_keys + unspecified_keys + specified_keys
        return {key: all_results[key] for key in new_key_order}

    def save_and_format_excel(self, all_results: Dict) -> str:
        if not all_results:
            return ""
        all_results = self._remove_columns(all_results)
        listModuleOID = self.get_column_data_fn('ModuleOID', deduplicate=True)
        all_results = self._reorder_dict_unspecified_first(all_results, listModuleOID)

        self.output_path.parent.mkdir(parents=True, exist_ok=True)
        with pd.ExcelWriter(self.output_path, engine='openpyxl', mode='w') as writer:
            for sheet_name, df in all_results.items():
                if not isinstance(df, pd.DataFrame) or df.empty:
                    continue
                clean_name = self._clean_sheet_name(sheet_name)
                df_str = df.astype(str).replace('nan', '', regex=False)
                df_str.to_excel(writer, sheet_name=clean_name, index=False)

        wb = openpyxl.load_workbook(self.output_path)
        for sheet_name in wb.sheetnames:
            ws = wb[sheet_name]
            self._set_sheet_text_format(ws)
            self._enable_sheet_filter(ws)
            self._set_sheet_font(ws)
        wb.save(self.output_path)
        return str(self.output_path.resolve())